#%%-------------------------------------------
import numpy as np
from qutip import *
import matplotlib.pyplot as plt
import scipy.io as scio
import time


si = qeye(2)
sx = sigmax()
sy = sigmay()
sz = sigmaz()
sp = sigmap()
sm = sigmam()

sx_list = []
sy_list = []
sz_list = []
sp_list = []
sm_list = []

N = 14
for n in range(N):
    op_list = []
    for m in range(N):
        op_list.append(si)

    op_list[n] = sx
    sx_list.append(tensor(op_list))

    op_list[n] = sy
    sy_list.append(tensor(op_list))

    op_list[n] = sz
    sz_list.append(tensor(op_list))
    
    op_list[n] = sp
    sp_list.append(tensor(op_list))
    
    op_list[n] = sm
    sm_list.append(tensor(op_list))



#crosstalk coumplings
gc_data = scio.loadmat("20qDirectCouplingMatrixExp.mat")
# print(gc_data)
gc_mat = np.zeros((20, 20))
for i in range(19):
    for j in range(i+1, 20):
        try:
            gc_ij = gc_data["q%d-q%d" %(i+1, j+1)][0][0]
        except KeyError:
            continue
        gc_mat[i, j] = gc_ij

gc_mat = gc_mat/1000
# print(np.max(gc_mat))

#qubit-cavity couplings
gl = np.array([27.6, 27.4, 29.1, 27.6, 26.5, 29.2, 30.1, 24.1, 27.7 , 27.3, 26.9, 29.1, 27.4, 26.3, 26.5, 27.3, 29.0, 24.6, 27.5])/1000 #Ghz

wr = 5.51 #cavity frequency
cn = 3 #Truncated cavity levels, exact when cn=8

adag = create(cn)
a = adag.dag()


def xy_Hamil(ij_l, delta_l, cn=3, rt=False):
    H = 0.
    H += 2 * np.pi * wr * tensor([si for j in range(N)]+[adag*a])
    print("          gi      gj      delta    gc    J_real")
    ql = [] 
    for k in range(len(ij_l)):
        i, j = ij_l[k]
        wq = wr - delta_l[k]
        print("Q%d, Q%d: %.4f, %.4f, %.4f, %.4f, %.5f" %(i, j, gl[i], gl[j], wq-wr, gc_mat[i, j], gl[i]*gl[j]/(wq-wr)+gc_mat[i, j]))
        if not (i in ql):
            H += 2 * np.pi * (wq/2) * (tensor(sz_list[i], qeye(cn)))
            H += 2 * np.pi * gl[i] * (tensor(sp_list[i], a) + tensor(sm_list[i], adag))
            ql.append(i)

        if not (j in ql):
            H += 2 * np.pi * (wq/2) * (tensor(sz_list[j], qeye(cn))) 
            H += 2 * np.pi * gl[j] * (tensor(sp_list[j], a) + tensor(sm_list[j], adag))
            ql.append(j)

    print("crosstalk all: ")
    ql = np.sort(ql)
    for i in range(len(ql)-1):
        qi = ql[i]
        for j in range(i+1, len(ql)):
            qj = ql[j]
            if not (gc_mat[qi, qj] == 0):
                print("gc%d_%d: %.5f" %(qi, qj, gc_mat[qi, qj]))
            H += 2 * np.pi * gc_mat[qi, qj] * tensor(sp_list[qi]*sm_list[qj], qeye(cn))
            H += 2 * np.pi * gc_mat[qi, qj] * tensor(sp_list[qj]*sm_list[qi], qeye(cn))
    
    return H


def J_real(ij_sub, deltal):
    Jl = []
    for k in range(len(ij_sub)):
        i, j = ij_sub[k]
        J_ij = - gl[i]*gl[j] / deltal[k]
        J_ij += gc_mat[i, j]
        Jl.append(J_ij)
    Jl = np.abs(Jl)
    return Jl

#detune frequency for 3 layers
deltal0 = np.array([0.25, 0.34, 0.4, 0.45, 0.5, 0.55, 0.60])
ij_list0 = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13]])

deltal1 = np.array([0.45, 0.34, 0.4, 0.5, 0.25, 0.55, 0.60])
ij_list1 = np.array([[0, 9], [10, 11], [4, 12], [2, 5], [7, 13], [1, 6], [3, 8]])
# print(1./J_real(ij_list1, deltal1))

deltal2 =  np.array([0.60, 0.55,  0.34, 0.45, 0.5, 0.4 , 0.25])
ij_list2 = np.array([[0, 10], [2, 7], [9, 12], [1, 3], [5, 11], [4, 8], [6, 13]])
# print(1./J_real(ij_list2, deltal2))

deltal3 = np.array([0.60, 0.4,  0.45, 0.34, 0.25, 0.5, 0.55])
ij_list3 =  np.array([[0, 11], [5, 6], [9, 10], [2, 3], [7, 12], [1, 4], [8, 13]])
# print(1./J_real(ij_list3, deltal3))

ij_list = [ij_list1, ij_list2, ij_list3]
deltall = [deltal1, deltal2, deltal3]


layer = 0
ij_sub = ij_list[layer]
deltal = deltall[layer]

# # if N < 14:
# #     ij_list_all = np.concatenate((ij_list1, ij_list2, ij_list3))
# #     ij_sub = [ij for ij in ij_list_all if max(ij) < N] 
# #     ij_sub = ij_sub[:N//2]
# #     deltal = deltal[:N//2]

if N < 14:
    ij_sub = ij_list0[:N//2]

Jl = J_real(ij_sub, deltal)

print("period: ", 1./Jl)
Hxy = xy_Hamil(ij_sub, deltal, cn=cn)


#-------initialize states and observables------
basis_init = [basis(2, 0) for i in range(N)]
for i, j in ij_sub:
    basis_init[i] = basis(2, 0)
    basis_init[j] = basis(2, 1)

psi0 = tensor(basis_init)
psi1 = tensor(psi0, basis(cn, 0))
obszl = [tensor(sz_list[i], qeye(cn)) for i in range(N)]
obscavi = tensor([si for j in range(N)]+[adag*a])

#------------time evolution----------------
tau = 1./min(Jl)
tlist = np.linspace(0, tau, 1000)
res = mesolve(Hxy, psi1, tlist, e_ops = obszl)
zlt = res.expect
np.save("data/z_%dbit_layer%d_delta" %(N, layer) + ", ".join(str(i) for i in deltal), zlt)

#%%-----------------plot---------------------------
fig, axs = plt.subplots(int(N/2), 1, figsize=(5, 9))
zlt = np.load("data/z_%dbit_layer%d_delta" %(N, layer) + ", ".join(str(i) for i in deltal) + ".npy")

for k in range(len(ij_sub)):
    i, j = ij_sub[k]
    tau = 1./Jl[k]
    tlist_i = tlist[tlist<tau]

    axs[k].plot(tlist_i, zlt[i][:len(tlist_i)], label = "Q %d" %(i))
    axs[k].plot(tlist_i, zlt[j][:len(tlist_i)], label = "Q %d" %(j))
    axs[k].set_title(r"$\Delta=%.2f$ GHz, $J=%.2f$ MHz" %(deltal[k], Jl[k]*1000))

    axs[k].legend(loc=2)
    axs[k].set_ylabel(r"$\langle \sigma_z \rangle$")
    axs[k].plot(tlist_i, [1]*len(tlist_i), '--', c='k')
    axs[k].plot(tlist_i, [-1]*len(tlist_i), '--', c='k')

axs[-1].set_xlabel("t(ns)")
plt.tight_layout()
fig.savefig("fig/z_%dbit_layer%d_delta" %(N, layer) + ", ".join(str(i) for i in deltal) + ".pdf", dpi=400)

